Training Loop

train_loader, val_loader = random_split(dataset)
net = Network()
optim = Optimizer(net.parameters(), lr=0.1)
loss_fn = Loss()

for epoch in range(epochs):
	for x, y in train_loader:
		# Forward pass
		classes = net(x)
		
		# Backward pass
		optim.zero_grad()
		loss = loss_fn(classes, y)
		loss.backward()
		optim.step()